#include "common.h"
#include "kernel_operator.h"
#include "tensor.h"
#include <cstdint>

#ifdef __CCE_KT_TEST__
#include <iostream>
#define DTYPE_X half;
#define DTYPE_Y half;
#endif
namespace AscendC
{

constexpr float my_pi = 3.14159265358979323846;

constexpr float spence_A[8] = {
    4.65128586073990045278E-5, 7.31589045238094711071E-3, 1.33847639578309018650E-1, 8.79691311754530315341E-1,
    2.71149851196553469920E0,  4.25697156008121755724E0,  3.29771340985225106936E0,  1.00000000000000000126E0,
};

constexpr float spence_B[8] = {
    6.90990488912553276999E-4, 2.54043763932544379113E-2, 2.82974860602568089943E-1, 1.41172597751831069617E0,
    3.63800533345137075418E0,  5.03278880143316990390E0,  3.54771340985225096217E0,  9.99999999999999998740E-1,
};

template <typename _DT_X, typename _DT_Y> class Kernel
{
  public:
    TPipe pipe;
    int32_t last_size = 0;
    CommonTiling tilingData;

    uint16_t flag = 0;

    DefInTensor(X);
    DefOutTensor(Y);

    CompareHelper<_DT_X> compare_helper;
    float pi_206 = my_pi * my_pi / 6.0;

    DefBufVECCALC(WW);
    DefBufVECCALC(YY);
    DefBufVECCALC(ZZ);
    DefBufVECCALC(CMP_RET);
    DefBufVECCALC(CMP_RET2);
    DefBufVECCALC(CMP_RET_SUM);
    DefBufVECCALC(FLAG);
    DefBufVECCALC(FLAG2);
    DefBufVECCALC(TEMP_0);
    DefBufVECCALC(TEMP_0_5);
    DefBufVECCALC(TEMP_1);
    DefBufVECCALC(TEMP_1_5);
    DefBufVECCALC(TEMP_2);
    DefBufVECCALC(TEMP_PI_206);

    DEF_MY_TENSOR_TEMP;
    Tensor<TypeOf(X)> x;
    Tensor<TypeOf(X)> y;

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR x_addr, GM_ADDR y_addr, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(X, x_addr, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y_addr, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);

        // 初始化buf
        compare_helper.Init(tilingData, pipe);
        InitTBufBuffer(WW, tilingData.bufferSize * sizeof(TypeOf(X)));
        InitTBufBuffer(YY, tilingData.bufferSize * sizeof(TypeOf(X)));
        InitTBufBuffer(ZZ, tilingData.bufferSize * sizeof(TypeOf(X)));
        InitTBufBuffer(CMP_RET, tilingData.bufferSize * sizeof(TypeOf(X)));
        InitTBufBuffer(CMP_RET2, tilingData.bufferSize * sizeof(TypeOf(X)));
        InitTbufWithDefault(CMP_RET_SUM, tilingData.bufferSize, TypeOf(X)(0));
        InitTbufWithDefault(FLAG, tilingData.bufferSize, TypeOf(X)(0));
        InitTbufWithDefault(FLAG2, tilingData.bufferSize, TypeOf(X)(0));
        InitTbufWithDefault(TEMP_0, tilingData.bufferSize, TypeOf(X)(0));
        InitTbufWithDefault(TEMP_0_5, tilingData.bufferSize, TypeOf(X)(0.5));
        InitTbufWithDefault(TEMP_1, tilingData.bufferSize, TypeOf(X)(1));
        InitTbufWithDefault(TEMP_1_5, tilingData.bufferSize, TypeOf(X)(1.5));
        InitTbufWithDefault(TEMP_2, tilingData.bufferSize, TypeOf(X)(2));
        InitTbufWithDefault(TEMP_PI_206, tilingData.bufferSize, TypeOf(X)(pi_206));

        INIT_MY_TENSOR_TEMP_AND_GET(tilingData.bufferSize, TypeOf(X));
        INIT_MY_TENSOR_TEMP(x);
        INIT_MY_TENSOR_TEMP(y);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize, false);
            Compute(i, tilingData.bufferSize, false);
            CopyOut(i, tilingData.bufferSize, false);
        }
        CopyIn(finnal_progress, tilingData.finalCalcCount, true);
        Compute(finnal_progress, tilingData.finalCalcCount, true);
        CopyOut(finnal_progress, tilingData.finalCalcCount, true);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(X, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        QueAllocSimple(Y);
        TBUF_GET_MY_TENSOR_TEMP;
        MY_TENSOR_UPDATE(x, LTensorName(X), calcCount);
        MY_TENSOR_UPDATE(y, LTensorName(Y), calcCount);

        TBufGet(WW, TypeOf(X));
        TBufGet(TEMP_0_5, TypeOf(X));
        TBufGet(TEMP_1, TypeOf(X));
        TBufGet(TEMP_0, TypeOf(X));
        TBufGet(CMP_RET, uint8_t);
        DEFINE_MY_TENSOR(cmp_ret, uint8_t, TBufTensorName(CMP_RET), calcCount);

        x += 10;
        x += TBufTensorName(TEMP_0_5);
        y = x + 2;
        y *= 1.5;
        y -= TypeOf(X)(6);
        y = x / y;
        y = x - y;
        y = x.exp();
        y = x.max_(0);

        cmp_ret = x == TBufTensorName(TEMP_1);
        cmp_ret = x != TBufTensorName(TEMP_1);
        cmp_ret = x > TBufTensorName(TEMP_1);
        cmp_ret = x < TBufTensorName(TEMP_1);
        cmp_ret = x >= TBufTensorName(TEMP_1);
        cmp_ret = x <= TBufTensorName(TEMP_1);

        y.apxy(TBufTensorName(TEMP_1), 2);
        x = TypeOf(Y)(40000);
        y = y.select(y > x);
        EnQue(Y);
        QueFree(X);
    }
    /*__aicore__ inline void Compute2(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        DeQueSimple(X);
        QueAllocSimple(Y);
        compare_helper.SetCalcCount(calcCount);
        TBufGet(WW, TypeOf(X));
        TBufGet(YY, TypeOf(X));
        TBufGet(ZZ, TypeOf(X));
        TBufGet(TEMP_0, TypeOf(X));
        TBufGet(TEMP_0_5, TypeOf(X));
        TBufGet(TEMP_1, TypeOf(X));
        TBufGet(TEMP_1_5, TypeOf(X));
        TBufGet(TEMP_2, TypeOf(X));
        TBufGet(TEMP_PI_206, TypeOf(X));
        TBufGet(CMP_RET, uint8_t);
        TBufGet(CMP_RET2, uint8_t);
        TBufGet(CMP_RET_SUM, uint16_t);
        TBufGet(FLAG, TypeOf(X));
        TBufGet(FLAG2, TypeOf(X));
        auto uint16_cmp_ret = TBufTensorName(CMP_RET).ReinterpretCast<uint16_t>();
        auto uint16_cmp_ret2 = TBufTensorName(CMP_RET).ReinterpretCast<uint16_t>();
        auto x_cmp_ret = TBufTensorName(CMP_RET).ReinterpretCast<TypeOf(X)>();
        auto x_cmp_ret2 = TBufTensorName(CMP_RET2).ReinterpretCast<TypeOf(X)>();

        // WW = x < 0 ？ nan : 0
        compare_helper.Begin(TBufTensorName(WW));
        compare_helper.IfGraterEqual(LTensorName(X), TBufTensorName(TEMP_0));
        compare_helper.Then(TBufTensorName(TEMP_0));
        compare_helper.ElseScalar(TypeOf(X)(float(0.0) / float(0.0)));
        compare_helper.End();
        compare_helper.Or2Out(TBufTensorName(CMP_RET_SUM));

        // YY = x == 1 ? 0 : 0
        compare_helper.Begin(TBufTensorName(YY));
        compare_helper.IfNotEqual(LTensorName(X), TBufTensorName(TEMP_1));
        compare_helper.Then(TBufTensorName(TEMP_0));
        compare_helper.End();
        compare_helper.Or2Out(TBufTensorName(CMP_RET_SUM));

        // ZZ = x == 0 ? pi_206 : 0
        compare_helper.Begin(TBufTensorName(ZZ));
        compare_helper.IfNotEqual(LTensorName(X), TBufTensorName(TEMP_0));
        compare_helper.Then(TBufTensorName(TEMP_0));
        compare_helper.ElseScalar(TypeOf(X)(pi_206));
        compare_helper.End();
        compare_helper.Or2Out(TBufTensorName(CMP_RET_SUM));

        // 将上面的结果进行统计
        // Y = WW + YY + ZZ
        Add(TBufTensorName(WW), TBufTensorName(WW), TBufTensorName(YY), calcCount);
        Add(LTensorName(Y), TBufTensorName(WW), TBufTensorName(ZZ), calcCount);

        // WW = 1 / x
        Div(TBufTensorName(WW), TBufTensorName(TEMP_1), LTensorName(X), calcCount);
        // if (x > 2.0) {
        //     x = 1.0 / x;
        //     flag = 1;
        // }
        compare_helper.Begin(LTensorName(X));
        compare_helper.IfGraterThan(LTensorName(X), TBufTensorName(TEMP_2));
        compare_helper.Then(TBufTensorName(WW));
        compare_helper.End();
        compare_helper.UserRetSelect(TBufTensorName(FLAG), TBufTensorName(TEMP_1));
        // if (x > 1.5) {
        //     w = (1.0 / x) - 1.0;
        //     flag = 1;
        // } else if (x < 0.5) {
        //     w = -x;
        //     flag2 == 1;
        // } else {
        //     w = x - 1.0;
        // }
        // YY = WW - 1
        Sub(TBufTensorName(YY), TBufTensorName(WW), TBufTensorName(TEMP_1), calcCount);
        compare_helper.Begin(TBufTensorName(WW));
        compare_helper.IfGraterThan(LTensorName(X), TBufTensorName(TEMP_1_5));
        compare_helper.Then(TBufTensorName(YY));
        compare_helper.End();
        compare_helper.UserRetSelect(TBufTensorName(FLAG2), TBufTensorName(TEMP_1));
        compare_helper.Copy2Out(uint16_cmp_ret);

        // YY = -X
        Sub(TBufTensorName(YY), TBufTensorName(TEMP_0), LTensorName(X), calcCount);
        compare_helper.Begin(TBufTensorName(ZZ));
        compare_helper.IfLessTan(LTensorName(X), TBufTensorName(TEMP_0_5));
        compare_helper.Then(TBufTensorName(YY));
        compare_helper.End();
        compare_helper.Or2Out(uint16_cmp_ret);

        // 合并结果
        Add(TBufTensorName(WW), TBufTensorName(WW), TBufTensorName(ZZ), calcCount);
        // 合并flag
        Add(TBufTensorName(FLAG), TBufTensorName(FLAG), TBufTensorName(YY), calcCount);

        // YY = X - 1.0
        Sub(TBufTensorName(YY), LTensorName(X), TBufTensorName(TEMP_1), calcCount);
        Not(uint16_cmp_ret, uint16_cmp_ret, calcCount);
        compare_helper.Begin(TBufTensorName(ZZ), &TBufTensorName(CMP_RET));
        compare_helper.Then(TBufTensorName(YY));
        compare_helper.End();

        // 合并结果
        Add(TBufTensorName(WW), TBufTensorName(WW), TBufTensorName(ZZ), calcCount);

        // y = -w * polevl(w, detail::spence_A, 7) / polevl(w, detail::spence_B, 7);
        // YY = -w
        Add(TBufTensorName(YY), TBufTensorName(TEMP_0), TBufTensorName(WW), calcCount);
        // ZZ = polevl(w, detail::spence_A, 7)
        polevl(TBufTensorName(WW), TBufTensorName(ZZ), spence_A, 7, calcCount);
        // x_cmp_ret = polevl(w, detail::spence_B, 7)
        polevl(TBufTensorName(WW), x_cmp_ret, spence_B, 7, calcCount);
        //  yy = -w * ZZ
        Mul(TBufTensorName(YY), TBufTensorName(YY), TBufTensorName(ZZ), calcCount);
        // yy = yy * x_cmp_ret
        Mul(TBufTensorName(YY), TBufTensorName(YY), x_cmp_ret, calcCount);

        compare_helper.Begin(x_cmp_ret);
        compare_helper.IfEqual(TBufTensorName(FLAG), TBufTensorName(TEMP_1));
        compare_helper.Then(TBufTensorName(YY));
        compare_helper.End();

        Ln(TBufTensorName(ZZ), LTensorName(X), calcCount);                          // z = std::log(x);
        Mul(TBufTensorName(WW), TBufTensorName(ZZ), TBufTensorName(ZZ), calcCount); // WW = ZZ * ZZ
        Muls(TBufTensorName(ZZ), TBufTensorName(WW), TypeOf(X)(-0.5), calcCount);   // ZZ = -0.5 * WW
        Sub(TBufTensorName(WW), TBufTensorName(ZZ), TBufTensorName(YY), calcCount); // WW = ZZ - YY

        compare_helper.Begin(TBufTensorName(ZZ));
        compare_helper.IfEqual(TBufTensorName(FLAG2), TBufTensorName(TEMP_1));
        compare_helper.Then(TBufTensorName(WW));
        compare_helper.End();

        Add(LTensorName(Y), LTensorName(Y), TBufTensorName(ZZ), calcCount);
        Add(LTensorName(Y), LTensorName(Y), x_cmp_ret, calcCount);

        EnQue(Y);
        QueFree(X);
    }*/
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount, bool is_final)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount);
    }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void spence(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    AscendC::Kernel<DTYPE_X, DTYPE_Y> op;
    op.Init(x, y, commonTiling);
    op.Process();
}